import sys, os
from telethon import TelegramClient, events
from telethon.tl.types import DocumentAttributeVideo, DocumentAttributeAudio
from telethon.sessions import StringSession
from urllib import request
import traceback
import subprocess
import m3u8
import asyncio
import typing
import ffmpeg


api_id = int(os.environ['API_ID'])
api_hash = os.environ['API_HASH']

CHAT_WITH_BOT_ID = os.environ['CHAT_WITH_BOT_ID']

client = TelegramClient(StringSession(os.environ['CLIENT_SESSION']), api_id, api_hash)

vid_format = '(best[ext=mp4,height<=1080]+best[ext=mp4,height<=480])[protocol^=http]/best[ext=mp4,height<=1080]+best[ext=mp4,height<=480]/best[ext=mp4,height<=1080]/best[ext=mp4]+worst[ext=mp4]/best[ext=mp4]/bestvideo[ext=mp4,height<=1080]+bestaudio[ext=m4a]/bestvideo[ext=mp4]+bestaudio[ext=m4a]/best'
worst_video_format = '(best[ext=mp4,height<=360])[protocol^=http]/best[ext=mp4,height<=360]/bestvideo[ext=mp4,height<=360]+bestaudio[ext=m4a]/best'
audio_format = 'bestaudio[ext=m4a]/bestaudio/best[ext=mp4,height<=480]/best[ext=mp4]/best'

TG_MAX_FILE_SIZE = 1500


class DumbReader(typing.BinaryIO):
    def write(self, s: typing.Union[bytes, bytearray]) -> int:
        pass
    def mode(self) -> str:
        pass
    def name(self) -> str:
        pass
    def close(self) -> None:
        pass
    def closed(self) -> bool:
        pass
    def fileno(self) -> int:
        pass
    def flush(self) -> None:
        pass
    def isatty(self) -> bool:
        pass
    def readable(self) -> bool:
        pass
    def readline(self, limit: int = -1) -> typing.AnyStr:
        pass
    def readlines(self, hint: int = -1) -> typing.List[typing.AnyStr]:
        pass
    def seek(self, offset: int, whence: int = 0) -> int:
        pass
    def seekable(self) -> bool:
        pass
    def tell(self) -> int:
        pass
    def truncate(self, size: int = None) -> int:
        pass
    def writable(self) -> bool:
        pass
    def write(self, s: typing.AnyStr) -> int:
        pass
    def writelines(self, lines: typing.List[typing.AnyStr]) -> None:
        pass
    def __enter__(self) -> 'typing.IO[typing.AnyStr]':
        pass
    def __exit__(self, type, value, traceback) -> None:
        pass

class FFMpegAV(DumbReader):

    def __init__(self, vformat, aformat=None, audio_only=False):
        _finput = ffmpeg.input(vformat['url'], **{"user-agent": user_agent, "loglevel": "error"})
        _fstream = None
        self.format = None
        if audio_only:
            self.format = 'mp3'
            acodec = None
            if 'acodec' in  vformat and vformat['acodec'] is not None:
                # if vformat['acodec'].startswith('mp4a'):
                #     acodec = 'm4a'
                if vformat['acodec'].startswith('mp3'):
                    acodec = 'mp3'

                if acodec != None:
                    _fstream = _finput.output('pipe:',
                                              format=acodec,
                                              acodec='copy',
                                              **{'vn': None})
                else:
                    _fstream = _finput.output('pipe:',
                                              format='mp3',
                                              acodec='mp3',
                                              **{'vn': None})
            else:
                _fstream = _finput.output('pipe:',
                                          format='mp3',
                                          acodec='mp3',
                                          **{'vn': None})
        else:
            _fstream = _finput.output('pipe:',
                                      format='mp4',
                                      vcodec='copy',
                                      acodec='copy',
                                      movflags='frag_keyframe')
        if aformat:
            self.stream = _fstream.global_args('-user-agent', user_agent, '-i', aformat['url'], '-map', '0:v', '-map', '1:a').run_async(pipe_stdout=True)
        else:
            self.stream = _fstream.run_async(pipe_stdout=True)

    def read(self, n: int = -1):
        return self.stream.stdout.read(n)

    def close(self) -> None:
        print('last data ', len(self.stream.stdout.read()))
        self.stream.kill()


def av_info(url, use_m3u8=False, audio_info=False):
    if use_m3u8:
        m3u8_obj = m3u8.load(url)
        url = m3u8_obj.segments[0].absolute_uri  # override for mediainfo call
        dur = 0
        for s in m3u8_obj.segments:
            if hasattr(s, 'duration'):
                dur += s.duration

    mediainf_args = None
    if audio_info:
        mediainf_args = '--Inform=Audio;%Duration%'
    else:
        mediainf_args = '--Inform=Video;%Width%\\n%Height%\\n%Duration%'
    mi_proc = subprocess.Popen(['mediainfo', mediainf_args, '2>', '/dev/null', url],
                             stdout=subprocess.PIPE,
                             stderr=subprocess.STDOUT)
    out = mi_proc.stdout.read()
    out = out.split(b'\n')

    w = h = None
    if not audio_info:
        w = int(out[0])
        h = int(out[1])
    if use_m3u8:
        dur = int(dur)
    else:
        dur = int(int(float(out[2 if not audio_info else 0]))/1000)

    if audio_info:
        return dur
    else:
        return w, h, dur

def video_format(url):
    mi_proc = subprocess.Popen(['mediainfo', '--Inform=General;%Format%', '2>', '/dev/null', url],
                               stdout=subprocess.PIPE,
                               stderr=subprocess.STDOUT)
    out = mi_proc.stdout.read()
    return out.split(b'\n')[0].decode('utf-8')

user_agent = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/74.0.3719.5 Safari/537.36'

def video_size(url, http_headers=None):
    size = 0
    head_req = request.Request(url, method='HEAD', headers=http_headers)
    with request.urlopen(head_req) as resp:
        size = int(resp.headers['Content-Length'])
    if size != 0:
        return size
    get_req = request.Request(url, method='GET', headers=http_headers)
    with request.urlopen(get_req) as resp:
        size = int(resp.headers['Content-Length'])
    
    return size



def m3u8_video_size(url, http_headers=None):
    m3u8_obj = m3u8.load(url)
    size = 0

    for seg in m3u8_obj.segments:
        head_req = request.Request(seg.absolute_uri, method='HEAD', headers=http_headers)
        with request.urlopen(head_req) as resp:
            size += int(resp.headers['Content-Length'])

    return size

async def main():
    chat_and_message_id = str(sys.argv[1])
    urls = str(sys.argv[2]).split(" ")
    mode = None if len(sys.argv) <= 3 else sys.argv[3]

    y_format = None
    playlist_start = None
    playlist_end = None
    # p - playlist video; pa - playlist audio; pw - playlist worse video
    if mode is not None:
        if mode.startswith('p') or mode.startswith('pa') or mode.startswith('pw'):
            pmode, prange = mode.split(':')
            _start, _end = prange.split('-')
            playlist_start = int(_start)
            playlist_end = int(_end)
            # cut "p" from mode variable if mode == "pa" or "pw"
            mode = pmode if len(pmode) == 1 else pmode[-1]

        if mode == 'a':
            # audio mode
            y_format = audio_format
        elif mode == 'w':
            # wordst video mode
            y_format = worst_video_format
        else:
            # normal mode
            y_format = vid_format
    else:
        y_format = vid_format

    for u in urls:
        try:
            params = {'format': y_format, 'noplaylist': True, 'youtube_include_dash_manifest': False}
            if playlist_start != None and playlist_end != None:
                if playlist_start == 0 and playlist_end == 0:
                    params['playliststart'] = 1
                    params['playlistend'] = 10
                else:
                    params['playliststart'] = playlist_start
                    params['playlistend'] = playlist_end
            else:
                params['playlist_items'] = '1'

            y = ydl.YoutubeDL(params)
            vinfo = y.extract_info(u, download=False)
        except Exception as e:
            if "Please log in or sign up to view this video" in e.__str__():
                if 'vk.com' in u or 'facebook.com' in u:
                    try:
                        params['username'] = os.environ['VIDEO_ACCOUNT_USERNAME']
                        params['password'] = os.environ['VIDEO_ACCOUNT_PASSWORD']
                        yy = ydl.YoutubeDL(params)
                        vinfo = yy.extract_info(u, download=False)
                    except Exception as e:
                        await client.send_message(CHAT_WITH_BOT_ID, chat_and_message_id+" "+e.__str__())
                        continue
                else:
                    await client.send_message(CHAT_WITH_BOT_ID, chat_and_message_id+" "+e.__str__())
                    continue
            elif 'are video-only' in e.__str__():
                try:
                    params['format'] = 'bestvideo[ext=mp4]'
                    yy = ydl.YoutubeDL(params)
                    vinfo = yy.extract_info(u, download=False)
                except Exception as e:
                    await client.send_message(CHAT_WITH_BOT_ID, chat_and_message_id+" "+e.__str__())
                    continue
            else:
                await client.send_message(CHAT_WITH_BOT_ID, chat_and_message_id+" "+e.__str__())
                traceback.print_exc()
                continue

        entries = None
        if '_type' in vinfo and vinfo['_type'] == 'playlist':
            entries = vinfo['entries']
        else:
            entries = [vinfo]

        for entry in entries:
            formats = entry.get('requested_formats')
            file_size = None
            chosen_format = None
            ffmpeg_av = None
            http_headers = None
            if 'http_headers' not in entry:
                if len(formats) > 0 and 'http_headers' in formats[0]:
                        http_headers = formats[0]['http_headers']
            else:
                http_headers = entry['http_headers']

            if formats is not None:
                for i, f in enumerate(formats):
                    if f['protocol'] in ['rtsp', 'rtmp', 'rtmpe', 'mms', 'f4m', 'ism', 'http_dash_segments']:
                        continue
                    if 'm3u8' in f['protocol']:
                        file_size = m3u8_video_size(f['url'], http_headers)
                    else:
                        file_size = video_size(f['url'], http_headers)

                    if f['protocol'] == 'https' and (True if ('acodec' in f and (f['acodec'] == 'none' or f['acodec'] == None)) else False):
                        # Dash video
                        vformat = f
                        mformat = None
                        vsize = video_size(vformat['url'], http_headers)
                        msize = 0
                        if len(formats) > i+1:
                            mformat = formats[i+1]
                            video_size(mformat['url'], http_headers)
                        file_size = vsize + msize + 10*1024*1024
                        if file_size/(1024*1024) < TG_MAX_FILE_SIZE:
                            ffmpeg_av = FFMpegAV(vformat, mformat)
                            chosen_format = f
                        break
                    if ('m3u8' in f['protocol'] and file_size / (1024*1024) <= TG_MAX_FILE_SIZE):
                        chosen_format = f
                        ffmpeg_av = FFMpegAV(chosen_format, audio_only=True if mode == 'a' else False)
                        break
                    if file_size / (1024 * 1024) <= TG_MAX_FILE_SIZE:
                        chosen_format = f
                        if mode == 'a' and not (chosen_format['ext'] == 'mp3'): #or chosen_format['acodec'].startswith('mp4a')):
                            ffmpeg_av = FFMpegAV(chosen_format, audio_only=True)
                        break

            else:
                if entry['protocol'] in ['rtsp', 'rtmp', 'rtmpe', 'mms', 'f4m', 'ism', 'http_dash_segments']:
                    await client.send_message(CHAT_WITH_BOT_ID, chat_and_message_id+" "+"ERROR: Failed find suitable video format")
                    return
                if 'm3u8' in entry['protocol']:
                    file_size = m3u8_video_size(entry['url'], http_headers)
                else:
                    file_size = video_size(entry['url'], http_headers)
                if ('m3u8' in entry['protocol'] and file_size / (1024*1024) <= TG_MAX_FILE_SIZE):
                    chosen_format = entry
                    ffmpeg_av = FFMpegAV(chosen_format, audio_only=True if mode == 'a' else False)
                if (file_size / (1024 * 1024) <= TG_MAX_FILE_SIZE):
                    chosen_format = entry
                    if mode == 'a' and not (chosen_format['ext'] == 'mp3'): #or chosen_format['acodec'].startswith('mp4a')):
                        ffmpeg_av = FFMpegAV(chosen_format, audio_only=True)


            try:
                if chosen_format is None and ffmpeg_av is None:
                    await client.send_message(CHAT_WITH_BOT_ID, chat_and_message_id+" "+"ERROR: Failed find suitable video format")
                    return
                if chosen_format['ext'] == 'unknown_video':
                    format = video_format(chosen_format['url'])
                    if format == 'MPEG-4':
                        chosen_format['ext'] = 'mp4'
                    else:
                        await client.send_message(CHAT_WITH_BOT_ID, chat_and_message_id+" "+"ERROR: Failed find suitable video format")
                        return
                if mode == 'a':
                    # we don't know real size due to converting formats
                    # so increase it in case of real size is less large then estimated
                    file_size += 200000
                file = await client.upload_file(ffmpeg_av if ffmpeg_av is not None else chosen_format['url'],
                                                file_name=entry['title'] + '.' + (chosen_format['ext'] if ffmpeg_av is None or ffmpeg_av.format is None else ffmpeg_av.format),
                                                file_size=file_size, http_headers=http_headers)

                width = height = duration = None
                if mode == 'a':
                    if ('duration' not in entry and 'duration' not in chosen_format):
                        duration = av_info(chosen_format['url'], use_m3u8=('m3u8' in chosen_format['protocol']), audio_info=True)
                    else:
                        duration = int(entry['duration']) if 'duration' not in entry else int(entry['duration'])

                elif ('duration' not in entry and 'duration' not in chosen_format) or ('width' not in chosen_format) or ('height' not in chosen_format):
                    width, height, duration = av_info(chosen_format['url'], use_m3u8=('m3u8' in chosen_format['protocol']))
                else:
                    width, height, duration = chosen_format['width'], chosen_format['height'], \
                                              int(entry['duration']) if 'duration' not in entry else int(entry['duration'])

                if ffmpeg_av is not None:
                    ffmpeg_av.close()

                attributes = None
                if mode == 'a':
                    attributes = DocumentAttributeAudio(duration, title=entry['title'])
                else:
                    attributes = DocumentAttributeVideo(duration,
                                                        width,
                                                        height,
                                                        supports_streaming=False if ffmpeg_av is not None else True)
                force_document = False
                if ffmpeg_av is None and (chosen_format['ext'] != 'mp4' and mode != 'a'):
                        force_document = True
                await client.send_file(CHAT_WITH_BOT_ID, file, video_note=False if mode == 'a' or force_document else True,
                                       voice_note= True if mode == 'a' else False,
                                       attributes=((attributes,) if not force_document else None), caption=chat_and_message_id, force_document=force_document)
            except Exception as e:
                print(e)
                traceback.print_exc()


client.start()

loop = asyncio.get_event_loop()
loop.run_until_complete(main())

